import tensorflow as tf
import numpy as np
import logging

from PyQt5.QtCore import pyqtSignal, QThread

class CustomCallback(tf.keras.callbacks.Callback):
    def __init__(self, trainTread):
        super(CustomCallback, self).__init__()
        self.trainTread = trainTread
        self.loss = ""
        self.accuracy = ""

    def on_train_batch_end(self, batch, logs={}):

        # print('batch：' + str(batch))
        loss = str(np.around(logs.get('loss'), 3))
        # print('loss 损失值：' + loss)
        accuracy = str(np.around(logs.get('accuracy'), 3))
        # print('accuracy 准确率：' + accuracy)

        self.loss = loss
        self.accuracy = accuracy

        self.trainTread.signal.emit("", loss, accuracy, "")


    def on_epoch_end(self, batch, logs={}):
        print('batch_size' + str(batch))
        self.trainTread.signal.emit("", self.loss, self.accuracy, 5 + (batch+1) * 3)


# 训练模型的线程
class TrainThread(QThread):
    # 传递string类型的参数
    signal = pyqtSignal(str, str, str, int)

    def __init__(self, ui):
        super(TrainThread, self).__init__()
        self.ui = ui

    def run(self):
        # 3. 训练
        callback = CustomCallback(self)
        history = self.ui.model.fit(self.ui.train_dataset, validation_data=self.ui.validate_dataset, epochs=self.ui.epochs, callbacks=[callback])
        self.signal.emit(str("OK"), callback.loss, callback.accuracy, 95)
